import matplotlib.pyplot as plt
import numpy as np
import sys,matplotlib
import matplotlib

filename = '/home/kjs/gadams_sample_mars_one/kinmars_changes_dust.out-100'
infile = open(filename, 'r')
temp = infile.readline()
ntime = 1000
nz = 118
dat = {}
runtime = 0
for i in range(ntime):
    print(i)
    while not 'TSTEP' in temp:
        temp = infile.readline()
    temp = temp.split()
    runtime +=float(temp[3])
    dat[runtime] = {}
    temp = infile.readline()
    while not 'MIXING' in temp:
        temp = infile.readline()
    alt = np.zeros(nz)
    blank = infile.readline()
    header = infile.readline().split()
    alt = np.zeros(nz)
    while len(header)>1:
        for ix in range(1,len(header)):
            dat[runtime][header[ix]]= np.zeros(nz)
        for iz in range(0,nz):
            temp = infile.readline().split()
            for ix in range(1,len(header)):
                dat[runtime][header[ix]][iz] = float(temp[ix+1])
            alt[iz] = float(temp[1])
        blank = infile.readline()
        header = infile.readline().split()
#for ipnames in range(0,len(speciesnames2)):
#    ax.plot(dat[files[1]][specienames2[ipnames]]/dat[files[0]][specienames2[ipnames]],alt,color=pclr[ipnames],linestyle=lns[fctr],linewidth=lw,label=specienames2[ipnames])
specienames = list(dat[1e4].keys())
dat2 = {}
time = list(dat.keys())
for i in specienames:
    dat2[i]=np.zeros((ntime,nz))
    for j in range(0,ntime):
        dat2[i][j,:]=dat[time[j]][i]
fix,ax = plt.subplots(ncols=1,nrows=4, figsize=(12,6)) # fix this
specienames2 = ['O2', 'O3', 'CO', 'HNO3', 'H2','H']
colors_of_species = ['red', 'blue', 'cyan', 'black', 'purple','green']
alt_slices = [2, 20, 80, 116]
ictr = 0
for i in specienames2:
    jctr = 0
    for j in alt_slices:
        #ax[jctr].loglog(time,dat2[i][:,j],colors_of_species[ictr],label=i)
        ax[jctr].semilogx(time,dat2[i][:,j]/dat2[i][0,j],colors_of_species[ictr],label=i)
        jctr +=1
    ictr +=1
# x = time
# y = mixing ratios
#species to plot
ax[0].set_title('Time Step Dust Profile')
ax[-1].set_xlabel('Time') #,fontsize=fontsize,fontname=fontname)
ax[0].set_ylabel('Mixing Ratio') #fontsize=fontsize,fontname=fontname)
ax[0].legend()
#ax.plot(time, dat2[i])
plt.savefig('timestep_dust.png')
